from skimage.feature import hog
from skimage import exposure
import matplotlib.pyplot as plt
from skimage import io

class HOGFeatureExtractor:
    def __init__(self, orientations=9, pixels_per_cell=(8, 8), cells_per_block=(2, 2)):
        """
        初始化HOG特征提取器的参数
        
        :param orientations: 梯度方向的数量，通常是9（表示0到180度分为9个区间）
        :param pixels_per_cell: 每个单元的像素大小，常见为(8, 8)像素
        :param cells_per_block: 每个块内包含的单元数，通常为(2, 2)
        """
        self.orientations = orientations
        self.pixels_per_cell = pixels_per_cell
        self.cells_per_block = cells_per_block

    def get_hog_feature(self, image):
        """
        从图像中提取HOG特征
        
        :param image: array-like，输入图像
        :return: 返回HOG特征向量fd和HOG图像hog_image
        """

        
        # 计算HOG特征
        fd, hog_image = hog(image, 
                            orientations=self.orientations, 
                            pixels_per_cell=self.pixels_per_cell, 
                            cells_per_block=self.cells_per_block, 
                            visualize=True)
        
        return fd, hog_image

    def visual_hog_feature(self, hog_feature, save_path=None):
        """
        可视化HOG特征图像
        
        :param hog_feature: HOG图像（hog_image）
        :param save_path: 如果指定，将保存HOG图像到给定路径
        """
        # 增强HOG图像的可视化效果
        hog_image_rescaled = exposure.rescale_intensity(hog_feature, in_range=(0, 10))
        
        # 绘制HOG图像
        plt.imshow(hog_image_rescaled, cmap=plt.cm.gray)
        plt.title('HOG Visualization')
        plt.axis('off')
        
        # 如果给定save_path，保存图像
        if save_path:
            plt.savefig(save_path)
        else:
            # 显示图像
            plt.show()
if __name__ == '__main__':
    # 创建HOG特征提取器
    import os
    hog_extractor = HOGFeatureExtractor()
    image_name_list=[
                    "1243.jpg",
                    "1244.jpg",
                    "1245.jpg",
                    "1246.jpg",
                    "1247.jpg"
                ]
    save_dir='./experiments/hog_visualization'
    os.makedirs(save_dir, exist_ok=True)
    vascular_dir='../Dataset/infantImages/vascular'
    for image_name in image_name_list:
        # 提取HOG特征
        hog_feature, hog_image = hog_extractor.get_hog_feature(os.path.join(vascular_dir,image_name))
        
        # 可视化HOG特征
        hog_extractor.visual_hog_feature(hog_image, save_path=os.path.join(save_dir,image_name))
